import torch
import torch.nn as nn
import torch.nn.functional as F

class GlobalNet(nn.Module):
    def __init__(self ,encode_information , channels , size , nBand, num_class):
        super(GlobalNet, self).__init__()
        self.encode_information = encode_information
        self.channels = channels
        self.size =size
        self.num_class = num_class
        self.layer_module = None

        # self.stem = nn.Sequential(
        #     nn.Conv2d(nBand, 10, kernel_size = 1, stride = 1, padding = 0, bias=False),
        #     nn.BatchNorm2d(10),
        # )

        input_fc_channel = self.channels[-1]
        for position in range(len(size)):
            if size[position] != None:
                input_fc_size = size[position]
            else:
                break
        self.layer_module = nn.ModuleList(encode_information)
        self.global_pooling = nn.AdaptiveAvgPool2d(1)
        input_fc_size = 1
        self.linear = nn.Linear(int(input_fc_channel) * int(input_fc_size) * int(input_fc_size ), num_class)

    def forward(self , x):
        # x = self.stem(x)
        for layer in self.layer_module:
            x = layer(x)
        x = self.global_pooling(x)
        out = x.view(x.size(0), -1)
        out = self.linear(out)
        return out
